package com.ai.langchain4j.config;

import com.ai.langchain4j.memory.PersistentChatMemoryStore;
import com.ai.langchain4j.service.ToolsService;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.SystemMessage;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.V;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import reactor.core.publisher.Flux;

/**
 * @Author: suren@517na.com
 * @CreateTime: 2025-05-25
 * @Description: langchain4j 配置类
 */

@Configuration
public class Langchain4JConfig {
    @Bean
    public OllamaChatModel langchain4jOllamaChatModel() {

        return OllamaChatModel.builder()
                .baseUrl("http://localhost:11434")
                //.apiKey(System.getenv("ALI_AI_KEY"))
                .modelName("deepseek-r1:1.5b")
                .build();
    }

    public interface Assistant {
        String chat(String message);

        //流式响应
        TokenStream stream(String message);
    }

    @Bean
    public Assistant assistant(ChatLanguageModel qwenChatModel,
                               StreamingChatLanguageModel qwenStreamingChatModel) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemory(chatMemory)
                .build();
        return assistant;
    }

    public interface AssistantUnique {
        String chat(@MemoryId int memoryId, @UserMessage String userMessage);
        //流式响应

        @SystemMessage("""
                您是一个航空公司聊天客服，请以友好的、助人为乐且愉快的方式回复。
                您正在通过在线聊天系统与客户互动。
                请提供有关预定或取消预定消息之前，您必须始终从用户处获取以下信息：预定好、客户姓名。
                请讲中文。
                今天的日期是{{current_date}}
                """)
        TokenStream stream(@MemoryId int memoryId, @UserMessage String userMessage, @V("current_date") String currentDate);
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<TextSegment>();
    }

    @Bean
    public AssistantUnique assistantUnique(ChatLanguageModel qwenChatModel,
                                           StreamingChatLanguageModel qwenStreamingChatModel,
                                           PersistentChatMemoryStore persistentChatMemoryStore,
                                           ToolsService toolsService,
                                           EmbeddingStore<TextSegment> embeddingStore,
                                           QwenEmbeddingModel qwenEmbeddingModel) {
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(5) //最相似的5个结果
                .minScore(0.6) //只找相似度在0.6以上的内容
                .build();
        AssistantUnique assistantUnique = AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .tools(toolsService)
                .contentRetriever(contentRetriever)
                .chatMemoryProvider(memoryId ->
                        MessageWindowChatMemory.builder()
                                .maxMessages(10)
                                .id(memoryId)
                                //自定义消息存储
                                .chatMemoryStore(persistentChatMemoryStore)
                                .build())
                .build();
        return assistantUnique;
    }
}
